import { warn } from "cc";
import { SamplerFactory } from "../Sampler";
import { Guid } from "../Utils/Guid";
import { IncomingMessage } from "./InComingMessage";
import { SideChannel } from "./SideChannel";

export enum EnvironmentDataTypes {
    Float = 0,
    Sampler = 1
}

export enum SamplerType {
    /**
     * @en
     *  Samples a reset parameter from a uniform distribution.
     * 
     * @zh
     * 从均匀分布中采样重置参数。
     */
    Uniform = 0,

    /**
     * @en
     *  Samples a reset parameter from a Gaussian distribution.
     * 
     * @zh
     * 从高斯分布中采样重置参数。
     */
    Gaussian = 1,

    /**
     * @en
     * Samples a reset parameter from a MultiRangeUniform distribution.
     * 
     * @zh
     * 从多范围均匀分布中采样重置参数。
     */
    MultiRangeUniform = 2
}

/**
 * @en
 * A side channel that manages the environment parameter values from Python. Currently
 * limited to parameters of type float.
 * 
 * @zh
 * 一个侧通道，用于管理来自 Python 的环境参数值。目前仅限于浮点类型的参数。
 */
export class EnvironmentParametersChannel extends SideChannel {
    private _parameters: { [key: string]: () => number } = {};
    private _registeredActions: { [key: string]: (value: number) => void } = {};

    private readonly _envParamsId = '534c891e-810f-11ea-a9d0-822485860400';

    constructor() {
        super();
        this.channelId = new Guid(this._envParamsId);
    }

    onMessageReceived(msg: IncomingMessage): void {
        const key = msg.readString();
        const type = msg.readInt32();

        if (EnvironmentDataTypes.Float === type) {
            const value = msg.readFloat32();
            this._parameters[key] = () => value;
            const action = this._registeredActions[key];
            action?.(value);
        } else if (EnvironmentDataTypes.Sampler === type) {
            const seed = msg.readInt32();
            const samplerType = msg.readInt32();
            let sampler = () => 0.0;

            if (SamplerType.Uniform === samplerType) {
                const min = msg.readFloat32();
                const max = msg.readFloat32();
                sampler = SamplerFactory.createUniformSampler(min, max, seed);
            } else if (SamplerType.Gaussian === samplerType) {
                const mean = msg.readFloat32();
                const stddev = msg.readFloat32();
                sampler = SamplerFactory.createGaussianSampler(mean, stddev, seed);
            } else if (SamplerType.MultiRangeUniform === samplerType) {
                const intervals = msg.readFloatList();
                sampler = SamplerFactory.createMultiRangeUniformSampler(intervals, seed);
            } else {
                warn('EnvironmentParametersChannel received an unknown data type.');
            }

            this._parameters[key] = sampler;
        } else {
            warn('EnvironmentParametersChannel received an unknown data type.');
        }
    }

    public getWithDefault(key: string, defaultValue: number): number {
        const valueOut = this._parameters[key];
        const hasKey = valueOut !== undefined;
        return hasKey ? valueOut() : defaultValue;
    }


    /**
     * 
     * @en
     * Registers a callback for the associated parameter key. Will overwrite any existing
     * actions for this parameter key.
     * 
     * @zh
     * 为关联的参数键注册回调。将覆盖此参数键的任何现有操作。
     * 
     * 
     * @param key 
     * @param action 
     */
    public registerCallback(key: string, action: (value: number) => void): void {
        this._registeredActions[key] = action;
    }


    /**
     * @en
     *  Returns all parameter keys that have a registered value.
     * 
     * @zh
     * 返回所有已注册值的参数键。
     * @returns 
     */
    public listParameters(): string[] {
        return Object.keys(this._parameters);
    }
}